import pandas as pd
import matplotlib.pyplot as plt


class PortAnalysis:
    """
    分析类
        - 组合分析：累计收益率、夏普比率、最大回撤
    """

    def __init__(self, port_value: dict, benchmark: pd.DataFrame = None, risk_free_rate=0.03):
        self.port_value = pd.DataFrame.from_dict(port_value, columns=['port_value'], orient='index')
        self.risk_free_rate = risk_free_rate

        if benchmark is not None:
            # 业绩基准数据校验
            required_columns = ['trading', 'close']
            missing_columns = [col for col in required_columns if col not in benchmark.columns]
            assert not missing_columns, f"缺少必要的列: {missing_columns}"
            if 'cum_return' not in benchmark.columns:
                benchmark['cum_return'] = benchmark['close'].pct_change().add(1).cumprod() - 1
            benchmark.set_index('trading', inplace=True)
        self.benchmark = benchmark

        # 分析结果存储
        self.port_indicators = {}

    def port_analysis(self):
        # 计算累计收益率
        _cum_return = self.port_value.pct_change().add(1).cumprod() - 1
        self.port_indicators['cum_return'] = _cum_return.values[-1][0]
        _annual_return = (1+_cum_return.values[-1][0]) ** (252/len(self.port_value)) - 1
        self.port_indicators['annualized_return'] = _annual_return

        # 计算夏普比率
        _port = self.port_value.copy()
        _port['daily_return'] = _port['port_value'].pct_change().dropna()
        _port['excess_return'] = _port['daily_return'] - self.risk_free_rate

        mean_cum_return = _annual_return - self.risk_free_rate
        std_excess_return = _port['excess_return'].std()

        if std_excess_return == 0:
            return 0.0

        sharpe_ratio = mean_cum_return / std_excess_return
        self.port_indicators['sharp_ratio'] = sharpe_ratio

        # 计算最大回撤
        cum_max = _port['port_value'].cummax()
        drawdown = (_port['port_value'] - cum_max) / cum_max

        max_drawdown_idx = drawdown.idxmin()
        max_drawdown_value = abs(drawdown.min())

        peak_series = _port['port_value'][:max_drawdown_idx]
        peak_idx = peak_series.idxmax() if len(peak_series) > 0 else None

        recovery_idx = None
        if max_drawdown_idx is not None:
            post_drawdown = _port['port_value'][max_drawdown_idx:]
            # 查找净值重新回到回撤前最高点的日期
            pre_peak_value = _port['port_value'][peak_idx] if peak_idx is not None else _port['port_value'].iloc[0]
            recovery_series = post_drawdown[post_drawdown >= pre_peak_value]
            recovery_idx = recovery_series.index[0] if len(recovery_series) > 0 else None

        max_drawdown = {
            'max_drawdown': max_drawdown_value,
            'peak_date': peak_idx,
            'trough_date': max_drawdown_idx,
            'recovery_date': recovery_idx
        }
        self.port_indicators['max_drawdown'] = max_drawdown

    def plot_return(self, save_path=None):
        # 累计收益率曲线图
        fig, ax = plt.subplots(figsize=(8, 4))

        port_cum_return = self.port_value.pct_change().add(1).cumprod() - 1
        ax.plot(port_cum_return.index, port_cum_return.values, label='Portfolio Cumulative Return', color='blue')

        if self.benchmark is not None:
            benchmark_data = self.benchmark.copy()
            ax.plot(benchmark_data.index, benchmark_data['cum_return'],
                    label='Benchmark Cumulative Return', color='red')

        # 设置图形属性
        ax.set_xlabel('Date')
        ax.set_ylabel('Cumulative Return')
        ax.set_title('Portfolio vs Benchmark Cumulative Return')
        ax.legend()
        ax.grid(True)

        # 自动调整日期格式
        fig.autofmt_xdate()

        plt.tight_layout()
        plt.show()
        if save_path is not None:
            plt.savefig(f'{save_path}.png')
